package qmlt.dataset.filter;

import java.util.ArrayList;
import java.util.List;

import qmlt.dataset.utils.EntropyUtils;

public class InfoGainDiscretizingFilter extends DiscretizingFilter
{
    @Override
    public List<Float> getDividers(List<Float> values, List<Object> classes)
    {
        int n = values.size();
        for (int i = 0; i < n; ++i)
        {
            for (int j = i + 1; j < n; ++j)
            {
                if (values.get(i) > values.get(j))
                {
                    float tmpf = values.get(j);
                    values.set(j, values.get(i));
                    values.set(i, tmpf);
                    Object tmps = classes.get(j);
                    classes.set(j, classes.get(i));
                    classes.set(i, tmps);
                }
            }
        }

        // find candidate divide points. those point to the position AFTER divide point
        List<Integer> candidateDivides = new ArrayList<Integer>();

        for (int i = 0; i < n - 1; ++i)
        {
            if (!classes.get(i).equals(classes.get(i + 1)))
                candidateDivides.add(i + 1);
        }

        // find the best divide point
        int bestDivide = 0;
        float bestInfoGain = 0;
        float beforeEnt = EntropyUtils.calculateEntropy(classes);

        for (int j : candidateDivides)
        {
            List<Object> subclasses1 = classes.subList(0, j);
            List<Object> subclasses2 = classes.subList(j, n);
            float infoGain = EntropyUtils.calculateInfoGain(n, beforeEnt, subclasses1, subclasses2);

            if (infoGain > bestInfoGain)
            {
                bestDivide = j;
                bestInfoGain = infoGain;
            }
        }

        List<Float> divider = new ArrayList<Float>();
        if (bestInfoGain > 0)
        {
            float v1 = values.get(bestDivide - 1);
            float v2 = values.get(bestDivide);
            divider.add((v1 + v2) / 2);
        }
        return divider;
    }
}
